import asyncio
import uuid

from langchain_tavily import TavilySearch

from app.agent.react import ReActAgent
from app.config import config
from app.prompt.career_plan import SYSTEM_PROMOMPT, FIRST_QUESTION

thread_id = None
messages = []
user_input = None


def init_thread_params():
    global thread_id, messages, user_input
    thread_id = str(uuid.uuid4())
    messages = [{"role": "system", "content": SYSTEM_PROMOMPT}]
    user_input = None


async def main():
    # Create the agent
    search = TavilySearch(tavily_api_key=config.get("tavily_search", "api_key"), max_results=2)
    reactagent = ReActAgent(name="Career plan agent", tools=[search])

    # Use the agent
    global thread_id, messages, user_input
    init_thread_params()
    while True:
        if user_input is None:
            print("欢迎踏上职业规划之旅！我是您的AI助手。为了更精准地帮助您，咱们从几个小问题开始吧。\n" + FIRST_QUESTION)

        user_input = input("发送消息（输入‘quit’退出，输入‘new'开启新对话）: ")
        if user_input.lower() == "quit":
            break

        if user_input.lower() == "new":
            init_thread_params()
            continue

        user_message = [{"role": "user", "content": user_input}]

        async for step in reactagent.ask_tool(user_msgs=user_message, thread_id=thread_id):
            step["messages"][-1].pretty_print()


if __name__ == "__main__":
    asyncio.run(main())
